#!/usr/bin/python
# -*- coding:utf-8 -*-

import numpy as np
import cv2
from time import time
from sklearn import svm
import matplotlib.pyplot as plt



#颜色图
data_train_color = cv2.imread("img/Montage_00.jpg")
data_y_color = cv2.imread("img/Montage_00-y.jpg")

#灰度图
data_train_gray = cv2.cvtColor(data_train_color, cv2.COLOR_BGR2GRAY).reshape(-1, 1)
data_y_gray = cv2.cvtColor(data_y_color, cv2.COLOR_BGR2GRAY).reshape(-1, 1)

#处理y值为 1 -1
# y_train = data_y_gray.astype(np.int16)
data_y_gray[data_y_gray < 255] = 1
data_y_gray[data_y_gray > 254] = 0
print (np.sum(data_y_gray>0))

#转换train为float类型
# data_train = data_train_gray.astype(np.float32)
# print(data_train.shape)

data_test_color = cv2.imread("img/Montage_11.jpg")
data_test_gray = cv2.cvtColor(data_test_color, cv2.COLOR_BGR2GRAY).reshape(-1, 1)
# data_test = data_test_gray.astype(np.float32)



# 分类器
t1 = time()
print("开始训练: ",t1)
clf = svm.SVC(kernel='rbf', decision_function_shape='ovo')
# clf = svm.SVC(C=0.8, kernel='rbf', gamma=20, decision_function_shape='ovr')
clf.fit(data_train_gray, data_y_gray.ravel())
t2 = time()
print("训练耗时: ",t2-t1)
print("训练结束，开始预测: ",t2)
y_hat = clf.predict(data_test_gray)
t3 = time()
print("预测耗时: ",t3-t2)
print("预测结束: ",t3)
#p rint (y_hat)

y_hat[y_hat==1] = 255
saveTrain = np.array(y_hat).reshape(512,512);
cv2.imwrite("new_img.jpg", saveTrain)
plt.imshow(saveTrain)
plt.show()



# data_train = np.array(np.random.random((10,2)), dtype=np.float32)
# y_train = np.array([1,-1,-1,1,-1,1,-1,-1,-1,1], dtype=np.int8)
#
# data_test = np.array(np.random.random((3,2)), dtype=np.float32)
#
# svm = cv2.ml.SVM_create();
# svm.setType(cv2.ml.SVM_C_SVC)  # SVM类型
# svm.setKernel(cv2.ml.SVM_LINEAR) # 使用线性核 SVM_RBF
# # svm.setDegree(0.0)
# # svm.setGamma(0.0)
# # svm.setCoef0(0.0)
# # svm.setC(1.0)
# # svm.setNu(0.0)
# # svm.setP(0.2)
# # svm.setClassWeights(None)
# svm.setTermCriteria((cv2.TERM_CRITERIA_COUNT, 100, 1.e-06))
#
# svm.train(data_train, cv2.ml.ROW_SAMPLE, y_train)
# response = svm.predict(data_test)
#
# print (response)
#


##save image
#data_y_gray[data_y_gray==1] = 255
#saveTrain = np.array(data_y_gray).reshape(512,512);
#cv2.imwrite("new_img.jpg", saveTrain)
